from enum import Enum
import json
import struct
import socket
import copy
from typing import Any, Union, List
from pydantic import BaseModel
from utils import logger

# =====================  常量  ====================
HEARTBEAT_TIMEOUT = 60  # heartbeat timeout unit seond
HEARTBEAT_PER = 5  # The number of seconds at which a heartbeat is sent
HEARTBEAT_CHECK_PER = 10  # Check heartbeat seconds

# ==================== 消息类型 ====================
class MessageType(Enum):
    Register = 1
    Heartbeat = 2
    GetTask = 3
    Task = 4
    TaskResult = 5
    NewTasks = 6
    SlaveInfo = 7
    TaskInfo = 8
    Refuse = 9
    Allow = 10
    ControlCommand = 11  # 客户端控制
    Exit = 12  # 目标主机退出
    Exception = 13
    PutongPutong = 14

# ==================== 任务信息 ====================
class TaskStatus(Enum):
    Error = 1
    Waiting = 2
    Running = 3
    Finished = 4
class EvalTask(BaseModel):
    """opencompass后端执行的格式"""
    id: str
    model: str
    datasets: list
    retry: int = 3
    max_num_workers: int = 40
class Task(BaseModel):
    """Java后端传入的格式"""
    id: str
    dataset: dict[str, list]
    model: str
    
    def convert_eval_task(self) -> list[EvalTask]:
        eval_list = []
        for dataset_name, sub_names in self.dataset.items():
            eval_list.append(
                EvalTask(id=self.id, model=self.model, datasets=sub_names)
            )
        return eval_list
class TaskInfo(BaseModel):
    """服务器保留的格式"""
    id: str
    PID: int
    dataset_name: str
    model: str
    subset_list: list[str]
    status: TaskStatus
    details: dict = {}

    def serialize(self) -> dict[str, Any]:
        t = self.model_dump()
        t['status'] = (self.status.name, self.status.value)
        return t
    @staticmethod
    def deserialize(data: dict) -> 'TaskInfo':
        t = copy.deepcopy(data)
        t['status'] = t['status'][1]
        t = TaskInfo.model_validate(t)
        return t        
# ==================== 主机信息 ====================
class HostInfo:
    def __init__(self, name, total_memory=None, available_memory=None, addr=None, port=None, times=None,
                 used_memory=None, mem_percent=None, on_task=False, on_task_num=0, tasks_info=None):
        self.name = name
        self.addr = addr
        self.port = port
        self.times = times
        # =============内存信息==============
        self.total_memory = total_memory
        self.available_memory = available_memory
        self.used_memory = used_memory
        self.mem_percent = mem_percent
        # =============任务信息==============
        self.on_task = on_task
        self.on_task_num = on_task_num
        self.tasks_info: list[TaskInfo] = tasks_info if tasks_info is not None else []
    
    def serialize(self):
        # 创建一个字典来存储序列化后的数据
        serialized_data = {
            "name": self.name,
            "addr": self.addr,
            "port": self.port,
            "times": self.times,
            "total_memory": self.total_memory,
            "available_memory": self.available_memory,
            "used_memory": self.used_memory,
            "mem_percent": self.mem_percent,
            "on_task": self.on_task,
            "on_task_num": self.on_task_num,
            "tasks_info": [t.serialize() for t in self.tasks_info]  # 序列化 tasks_info
        }
        return serialized_data
    
    @staticmethod
    def deserialize(data: dict) -> 'HostInfo':
        if not isinstance(data, dict):
            raise ValueError("HostInfo.deserialize expects a dictionary")
        hostinfo = copy.deepcopy(data)
        new_task_list = []
        for t in hostinfo.get('tasks_info', []):
            try:
                new_task_list.append(TaskInfo.deserialize(t))
            except Exception as e:
                logger.error(f"Failed to deserialize TaskInfo: {e}")
                # 可以选择跳过有问题的任务，或者用默认值填充
                continue

        hostinfo['tasks_info'] = new_task_list
        return HostInfo(**hostinfo)
        
    def set_task_info(self, task_info: Union['TaskInfo', List['TaskInfo']]):
        if isinstance(task_info, TaskInfo):
            self.tasks_info.append(task_info)
            self.on_task = True
            self.on_task_num += 1
        elif isinstance(task_info, list) and all(isinstance(t, TaskInfo) for t in task_info):
            self.on_task_num += len(task_info)
            self.tasks_info.extend(task_info)
            self.on_task = True
        else:
            raise TypeError("task_info must be an instance of TaskInfo or a list of TaskInfo")
        
    def get_task_info(self):
        return self.tasks_info

# ================== 构建消息方法 ==================
def build_register(client_id: str) -> dict:
    return {
        'type': MessageType.Register,
        'client_id': client_id,
    }

def build_allow(data=None) -> dict:
    return {
        'type': MessageType.Allow,
        'data': data,
    }

def build_refuse(message: str) -> dict:
    return {
        'type': MessageType.Refuse,
        'message': message
    }

def build_heartbeat(client_id: str, data=None) -> dict:
    return {
        'type': MessageType.Heartbeat,
        'client_id': client_id,
        'data': data, # 存储slave主机信息
    }

def build_new_tasks() -> dict:
    return {
        'type': MessageType.NewTasks,
    }

def build_get_task() -> dict:
    return {
        'type': MessageType.GetTask,
    }

def build_task(task: Task) -> dict:
    return {
        'type': MessageType.Task,
        'task': task.model_dump() if task else None
    }

def build_task_result(task_id: int, result: dict) -> dict:
    """如果result是None, status为False"""
    return {
        'type': MessageType.TaskResult,
        'task_id': task_id,
        'data': result,
        'status': True if result else False,
    }

def build_exit() -> dict:
    return {
        'type': MessageType.Exit,
    }
    
def build_exception(message: str) -> dict:
    return {
        'type': MessageType.Exception,
        'message': message,
    }
    
def build_slave_info(data: dict=None) -> dict:
    return {
        'type': MessageType.SlaveInfo,
        'data': data,
    }
    
def build_putong_putong(data: dict=[]) -> dict:
    return {
        'type': MessageType.PutongPutong,
        'data': data,
    }

# ==================== 协议处理 ====================
class MessageProtocol:
    HEADER_SIZE = 4  # 4字节长度头
    
    @staticmethod
    def send(sock: socket.socket, data: dict):
        """发送带长度头的消息"""
        data['type'] = data['type'].value
        msg = json.dumps(data).encode('utf-8')
        header = struct.pack('!I', len(msg))  # 网络字节序
        sock.sendall(header + msg)
    
    @staticmethod
    def recv(sock: socket.socket) -> dict:
        """接收完整消息"""
        header = sock.recv(MessageProtocol.HEADER_SIZE)
        if not header:
            return None
        msg_len = struct.unpack('!I', header)[0]
        
        chunks = []
        bytes_recd = 0
        while bytes_recd < msg_len:
            chunk = sock.recv(min(msg_len - bytes_recd, 4096))
            if not chunk:
                raise ConnectionError("Connection closed")
            chunks.append(chunk)
            bytes_recd += len(chunk)
        data = json.loads(b''.join(chunks).decode('utf-8'))
        data['type'] = MessageType(data['type'])
        return data